"""
快速数据检查 - 打印随机样本的标签信息
"""

import numpy as np
import os

# 导入项目模块
import sys
sys.path.insert(0, os.path.dirname(os.path.dirname(os.path.abspath(__file__))))

from data_generator import SpectrumDataGenerator, SignalLabel


def print_label(label: SignalLabel, index: int):
    """格式化打印标签信息"""
    
    print(f"\n{'='*80}")
    print(f"样本 {index} - 标签信息")
    print(f"{'='*80}")
    
    # 谐波信息
    print(f"\n【谐波特征】")
    print(f"  检测: {'✓ 是' if label.has_harmonic else '✗ 否'}")
    if label.has_harmonic:
        print(f"  谐波簇数: {len(label.harmonic_clusters)}")
        for i, cluster in enumerate(label.harmonic_clusters):
            print(f"\n  簇 {i+1}:")
            print(f"    基频: {cluster['base_freq']:.2f} Hz")
            print(f"    谐波数: {cluster['num_harmonics']}")
            frequencies = [cluster['base_freq'] * (j+1) for j in range(cluster['num_harmonics'])]
            print(f"    频率: {[f'{f:.1f}' for f in frequencies]} Hz")
            print(f"    幅值: {[f'{a:.3f}' for a in cluster['amplitudes']]}")
    
    # 调制信息
    print(f"\n【调制特征】")
    print(f"  检测: {'✓ 是' if label.has_modulation else '✗ 否'}")
    if label.has_modulation:
        print(f"  调制数量: {len(label.modulation_info)}")
        for i, mod in enumerate(label.modulation_info):
            print(f"\n  调制 {i+1}:")
            print(f"    载波频率: {mod['carrier_freq']:.2f} Hz")
            print(f"    调制频率: {mod['modulation_freq']:.2f} Hz")
            print(f"    边带数量: {mod['num_sidebands']}")
            sideband_space = mod.get('sideband_spacing', mod['modulation_freq'])
            print(f"    边带间距: {sideband_space:.2f} Hz")
            print(f"    边带位置: ±{sideband_space:.2f} Hz around {mod['carrier_freq']:.2f} Hz")
    
    # 轴承故障
    print(f"\n【轴承故障】")
    print(f"  检测: {'✓ 是' if label.has_bearing_fault else '✗ 否'}")
    if label.has_bearing_fault and label.bearing_fault_info:
        info = label.bearing_fault_info
        print(f"  故障类型: {info.get('type', 'unknown')}")
        print(f"  故障频率: {info.get('fault_freq', 0):.2f} Hz")
        print(f"  共振频率: {info.get('resonance_freq', 0):.2f} Hz")
    
    print()


def main():
    """快速检查数据"""
    
    print("="*80)
    print("快速数据检查")
    print("="*80)
    
    # 创建生成器
    generator = SpectrumDataGenerator(fs=5120, n_samples=16384, seed=42)
    
    print(f"\n配置:")
    print(f"  - 采样频率: {generator.params.fs} Hz")
    print(f"  - 采样点数: {generator.params.n_samples}")
    print(f"  - 信号时长: {generator.params.duration:.4f} s")
    print(f"  - 频谱长度: {8192}")
    
    # 生成不同类型的数据
    print(f"\n生成不同类型的样本...")
    
    # 1. 纯谐波
    spectrum1, label1 = generator.generate_pure_harmonic()
    print_label(label1, 1)
    print(f"  频谱形状: {spectrum1.shape}")
    print(f"  峰值位置: {np.argmax(spectrum1) * (5120/16384):.2f} Hz")
    
    # 2. 调制信号
    spectrum2, label2 = generator.generate_modulation_signal()
    print_label(label2, 2)
    
    # 3. 谐波+调制
    spectrum3, label3 = generator.generate_harmonic_with_modulation()
    print_label(label3, 3)
    
    # 4. 轴承故障
    spectrum4, label4 = generator.generate_bearing_fault()
    print_label(label4, 4)
    
    # 5. 复合信号
    spectrum5, label5 = generator.generate_complex_signal()
    print_label(label5, 5)
    
    # 批量生成
    print(f"\n{'='*80}")
    print(f"批量生成测试")
    print(f"{'='*80}")
    
    batch_size = 5
    spectra, labels = generator.generate_batch(batch_size)
    
    print(f"\n批量信息:")
    print(f"  - 批量大小: {batch_size}")
    print(f"  - 频谱形状: {spectra.shape}")
    print(f"  - 标签类型: {type(labels).__name__}")
    
    print(f"\n标签分布:")
    harmonic_count = sum(1 for l in labels if l.has_harmonic)
    modulation_count = sum(1 for l in labels if l.has_modulation)
    bearing_count = sum(1 for l in labels if l.has_bearing_fault)
    complex_count = sum(1 for l in labels if l.has_harmonic and l.has_modulation)
    
    print(f"  - 含谐波: {harmonic_count}/{batch_size}")
    print(f"  - 含调制: {modulation_count}/{batch_size}")
    print(f"  - 含轴承故障: {bearing_count}/{batch_size}")
    print(f"  - 复合信号: {complex_count}/{batch_size}")
    
    print(f"\n✓ 数据检查完成！")
    print(f"\n提示: 运行 python temp_scripts/visualize_training_data.py 查看可视化")


if __name__ == "__main__":
    main()

